# yolopredict.py：
import tkinter as tk
from tkinter import ttk, filedialog
import os
# 假设 shared_data.py 在同目录下
from shared import SharedData


class PredictionParamsGUI:
    def __init__(self, root, shared_data, app):
        self.root = root
        self.app = app  # 保存主程序引用
        self.shared_data = shared_data  # 存储共享数据
        self.root.title("Ultralytics - 预测参数")
        self.root.geometry("600x400")
        self.root.resizable(True, True)

        # 创建 UI 组件
        self.create_widgets()

    def create_widgets(self):
        # ========== 主滚动区域：Canvas + Scrollbar ==========
        self.canvas = tk.Canvas(self.root)
        scrollbar = ttk.Scrollbar(self.root, orient="vertical", command=self.canvas.yview)
        self.scrollable_frame = tk.Frame(self.canvas)

        self.scrollable_frame.bind(
            "<Configure>",
            lambda e: self.canvas.configure(scrollregion=self.canvas.bbox("all"))
        )

        self.canvas.create_window((0, 0), window=self.scrollable_frame, anchor="nw")
        self.canvas.configure(yscrollcommand=scrollbar.set)

        self.canvas.pack(side="left", fill="both", expand=True)
        scrollbar.pack(side="right", fill="y")

        # ========== 允许鼠标滚轮滚动 ==========
        def _on_mousewheel(event):
            self.canvas.yview_scroll(int(-1 * (event.delta / 120)), "units")

        self.canvas.bind_all("<MouseWheel>", _on_mousewheel)

        def _bind_mousewheel(event):
            self.canvas.focus_set()

        self.canvas.bind("<Enter>", _bind_mousewheel)

        # ========== 居中容器 ==========
        main_frame = tk.Frame(self.scrollable_frame)
        main_frame.pack(expand=True, padx=20, pady=20)

        # 配置网格列权重（用于居中）
        for i in range(4):  # 0,1,2,3 列
            main_frame.grid_columnconfigure(i, weight=1)

        row_idx = 0

        # ========== 标题 ==========
        title_label = tk.Label(main_frame, text="数据源（图片、视频路径/摄像头编号）", font=("Arial", 10, "bold"), anchor='center')
        title_label.grid(row=row_idx, column=1, columnspan=2, sticky='ew', pady=(0, 5))

        row_idx += 1

        # ========== 数据源输入行 ==========
        src_frame = tk.Frame(main_frame)
        src_frame.grid(row=row_idx, column=1, columnspan=2, sticky='ew', pady=5)

        self.source_entry = tk.Entry(src_frame, width=35)
        source_text = self.get_current_source_text()
        self.source_entry.insert(0, source_text)
        self.source_entry.pack(side=tk.LEFT, padx=(0, 5))

        self.select_file_btn = tk.Button(src_frame, text="选择文件", command=self.open_file_dialog, width=8)
        self.select_file_btn.pack(side=tk.LEFT)

        row_idx += 1

        # ========== 预测参数区域 ==========
        param_frame = tk.Frame(main_frame)
        param_frame.grid(row=row_idx, column=1, columnspan=2, sticky='ew', pady=10)

        # 定义参数列表（按两列分组）
        params = [
            ("project", "项目目录", "string"),
            ("confidence", "置信度阈值", "float"),
            ("iou", "IoU 阈值", "float"),
            ("image_size", "输入尺寸", "int"),
            ("device", "设备选择", "choice"),
            ("show", "实时显示", "bool"),
            ("save_frames", "保存帧", "bool"),
            ("save_txt", "保存标签txt", "bool"),
            ("save_conf", "保存置信度", "bool"),
            ("save_crop", "保存裁剪目标", "bool"),
            ("show_labels", "显示标签", "bool"),
            ("show_conf", "显示置信度", "bool"),
            ("show_boxes", "显示边框", "bool")
        ]

        self.param_vars = {}

        for i, (key, label, param_type) in enumerate(params):
            col = i % 2  # 0 或 1
            row = i // 2  # 0~5 行

            if param_type == "string":
                var = tk.StringVar(value=self.shared_data.model_config.get(key, ""))
                lbl = tk.Label(param_frame, text=label, anchor='w', font=("Arial", 9))
                lbl.grid(row=row, column=col*2, sticky='w', padx=(0, 5), pady=3)
                entry = tk.Entry(param_frame, textvariable=var, width=12)
                entry.grid(row=row, column=col*2+1, padx=(0, 10), pady=3)
                self.param_vars[key] = var

            elif param_type in ("float", "int"):
                var = tk.StringVar(value=str(self.shared_data.model_config.get(key, "")))
                lbl = tk.Label(param_frame, text=label, anchor='w', font=("Arial", 9))
                lbl.grid(row=row, column=col*2, sticky='w', padx=(0, 5), pady=3)
                entry = tk.Entry(param_frame, textvariable=var, width=12)
                entry.grid(row=row, column=col*2+1, padx=(0, 10), pady=3)
                self.param_vars[key] = var

            elif param_type == "choice" and key == "device":
                var = tk.StringVar(value=self.shared_data.model_config.get(key, "cuda:0"))
                lbl = tk.Label(param_frame, text=label, anchor='w', font=("Arial", 9))
                lbl.grid(row=row, column=col*2, sticky='w', padx=(0, 5), pady=3)
                # 下拉框
                device_combo = ttk.Combobox(
                    param_frame,
                    textvariable=var,
                    values=["cpu", "cuda:0"],
                    state="readonly",
                    width=10
                )
                device_combo.grid(row=row, column=col*2+1, padx=(0, 10), pady=3)
                self.param_vars[key] = var

            else:  # bool
                var = tk.BooleanVar(value=self.shared_data.model_config.get(key, False))
                cb = tk.Checkbutton(param_frame, text=label, variable=var, font=("Arial", 9))
                cb.grid(row=row, column=col*2, columnspan=2, sticky='w', padx=5, pady=3)
                self.param_vars[key] = var

        row_idx += 1

        # ========== 按钮区域 ==========
        button_frame = tk.Frame(main_frame)
        button_frame.grid(row=row_idx, column=1, columnspan=2, pady=15)

        self.start_btn = tk.Button(button_frame, text="启动", command=self.start_inference, width=8)
        self.start_btn.grid(row=0, column=0, padx=5)

        self.pause_btn = tk.Button(button_frame, text="暂停", command=self.pause_inference, width=8)
        self.pause_btn.grid(row=0, column=1, padx=5)

        self.stop_btn = tk.Button(button_frame, text="停止", command=self.stop_inference, width=8)
        self.stop_btn.grid(row=0, column=2, padx=5)

        row_idx += 1

        # ========== 保存参数按钮 ==========
        self.save_params_btn = tk.Button(
            main_frame,
            text="保存参数",
            command=self.save_params_to_yaml,
            width=15,
            bg="#4CAF50",
            fg="white",
            font=("Arial", 10, "bold")
        )
        self.save_params_btn.grid(row=row_idx, column=1, columnspan=2, pady=10)

        
    def get_current_source_text(self):
        """根据 shared_data 获取当前源的文本表示"""
        if self.shared_data.current_source == "picture" and self.shared_data.image_paths:
            return self.shared_data.image_paths[0]
        elif self.shared_data.current_source == "video" and self.shared_data.video_paths:
            return self.shared_data.video_paths[0]
        elif self.shared_data.current_source == "camera":
            return "摄像头设备 - 0"
        else:
            return "默认摄像头 - 0"

    def open_file_dialog(self):
        """打开文件选择对话框，并更新 shared_data 中的源文件，同时在主界面显示"""
        file_path = filedialog.askopenfilename(
            title="选择图像或视频文件",
            filetypes=[
                ("Image files", "*.jpg *.jpeg *.png *.bmp *.tiff"),
                ("Video files", "*.mp4 *.avi *.mov *.wmv *.mkv"),
                ("All files", "*.*")
            ]
        )
        if not file_path:
            return  # 用户取消选择

        # 更新 Entry 显示
        self.source_entry.delete(0, tk.END)
        self.source_entry.insert(0, file_path)

        # === 更新 shared_data 中的源信息 ===
        _, ext = os.path.splitext(file_path.lower())
        image_exts = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'}
        video_exts = {'.mp4', '.avi', '.mov', '.wmv', '.mkv'}

        if ext in image_exts:
            self.shared_data.current_source = "picture"
            self.shared_data.image_paths = [file_path]
            self.shared_data.video_paths = []
            print(f"已设置图片源: {file_path}")

            # ✅ 主动在主界面显示图片
            from file_operations import cmd_open_image_by_path
            cmd_open_image_by_path(self.shared_data, self.app, file_path)

        elif ext in video_exts:
            self.shared_data.current_source = "video"
            self.shared_data.video_paths = [file_path]
            self.shared_data.image_paths = []
            print(f"已设置视频源: {file_path}")

            # ✅ 主动在主界面打开并播放视频
            from file_operations import cmd_open_video_by_path
            cmd_open_video_by_path(self.shared_data, self.app, file_path)

        else:
            print(f"警告：不支持的文件类型 {ext}")
            return

        # ✅ 更新主界面文件列表（如果存在）
        if "file_listbox" in self.shared_data.ui_elements:
            listbox = self.shared_data.ui_elements["file_listbox"]
            listbox.delete(0, tk.END)
            listbox.insert(tk.END, os.path.basename(file_path))

    def start_inference(self):
        """启动推理：更新配置 + 使用当前数据源路径进行预测"""
        source_text = self.source_entry.get().strip()
        if not source_text:
            print("❌ 未指定数据源路径")
            return

        # === 1. 更新 shared_data.model_config 中的参数 ===
        for key, var in self.param_vars.items():
            if isinstance(var, tk.StringVar):
                value = var.get().strip()
                if not value:
                    print(f"❌ 参数 '{key}' 不能为空")
                    return
                try:
                    if key in ["confidence", "iou"]:
                        self.shared_data.model_config[key] = float(value)
                    elif key == "image_size":
                        self.shared_data.model_config[key] = int(value)
                    else:
                        self.shared_data.model_config[key] = value
                except ValueError as e:
                    print(f"❌ 参数 '{key}' 格式错误: {e}")
                    return
            elif isinstance(var, tk.BooleanVar):
                self.shared_data.model_config[key] = var.get()

        # === 2. 判断源类型并更新 shared_data ===
        if "摄像头" in source_text or "camera" in source_text.lower():
            self.shared_data.current_source = "camera"
            self.shared_data.image_paths = []
            self.shared_data.video_paths = []
            print("✅ 已设置为摄像头源: 0")
        else:
            if not os.path.exists(source_text):
                print(f"❌ 文件不存在: {source_text}")
                return

            _, ext = os.path.splitext(source_text.lower())
            image_exts = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'}
            video_exts = {'.mp4', '.avi', '.mov', '.wmv', '.mkv'}

            if ext in image_exts:
                self.shared_data.current_source = "picture"
                self.shared_data.image_paths = [source_text]
                self.shared_data.video_paths = []
                print(f"✅ 已设置图片源: {source_text}")
            elif ext in video_exts:
                self.shared_data.current_source = "video"
                self.shared_data.video_paths = [source_text]
                self.shared_data.image_paths = []
                print(f"✅ 已设置视频源: {source_text}")
            else:
                print(f"❌ 不支持的文件类型: {ext}")
                return

        # === 3. 触发主程序开始推理 ===
        if hasattr(self.shared_data, 'on_start_inference'):
            self.shared_data.on_start_inference()
        else:
            print("⚠️ 未设置 on_start_inference 回调")

        print(f"✅ 已设置数据源和预测参数并启动推理")

    def pause_inference(self):
        """调用主程序的暂停/恢复推理函数"""
        from infer_operations import cmd_infer_pause
        cmd_infer_pause(self.app)

    def stop_inference(self):
        """调用主程序的停止推理函数"""
        from infer_operations import cmd_infer_stop
        cmd_infer_stop(self.app)

    def save_params_to_yaml(self):
        """将当前参数保存为 predict.yaml 文件"""
        import yaml

        # 收集当前参数（不包括 source）
        config = {}
        for key, var in self.param_vars.items():
            if isinstance(var, tk.StringVar):
                value = var.get().strip()
                # 类型转换
                if key in ["confidence", "iou"]:
                    try:
                        value = float(value)
                    except ValueError:
                        print(f"❌ 参数 '{key}' 格式错误，使用默认值")
                        value = self.shared_data.model_config.get(key, 0.0)
                elif key == "image_size":
                    try:
                        value = int(value)
                    except ValueError:
                        print(f"❌ 参数 '{key}' 格式错误，使用默认值")
                        value = self.shared_data.model_config.get(key, 640)
                config[key] = value
            elif isinstance(var, tk.BooleanVar):
                config[key] = var.get()

        # 写入 YAML 文件
        try:
            os.makedirs("cfg", exist_ok=True)
            yaml_path = "cfg/predict.yaml"
            with open(yaml_path, 'w', encoding='utf-8') as f:
                yaml.dump(config, f, allow_unicode=True, default_flow_style=False, sort_keys=False)
            print(f"✅ 参数已保存至: {yaml_path}")
            import tkinter.messagebox as msgbox
            msgbox.showinfo("成功", f"参数已保存至:\n{yaml_path}")
        except Exception as e:
            print(f"❌ 保存失败: {e}")
            import tkinter.messagebox as msgbox
            msgbox.showerror("错误", f"保存失败:\n{str(e)}")
